from server import Server

import time
import signal
import sys
import socket
import threading
import logging
from queue import Queue
from utils import *
from common import *

# ==================== 主服务端 ====================
class MasterServer:
    def __init__(self, host='0.0.0.0', port=9000):
        self.host = host
        self.port = port
        self.listen_socket = None
        self.clients = {}  # {client_id: (socket, addr, last_active)}
        self.clients_info = {}  # {client_id: {on_task: bool}}
        self.task_queue = Queue()
        self.lock = threading.Lock()
        self.running = False
        self.config = get_config("./config.cfg")
        
        # 日志配置
        logging.basicConfig(
            filename='master.log',
            format='%(asctime)s [%(levelname)s] %(message)s',
            level=logging.INFO
        )
        
    def start(self):
        """启动主服务"""
        self.running = True
        accept_connections_thread = threading.Thread(target=self._accept_connections, daemon=True)
        accept_connections_thread.start()
        # check_heartbeats_thread = threading.Thread(target=self._check_heartbeats, daemon=True)
        # check_heartbeats_thread.start()
        logging.info("Master server started")
        # check_heartbeats_thread.join()
        accept_connections_thread.join()
        logging.info("Master server stopped")

    def _accept_connections(self):
        """接受客户端连接"""
        with socket.socket() as s:
            linger = struct.pack('ii', 1, 5)  # 启用 SO_LINGER，延迟时间为 10 秒
            s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger)
            s.bind((self.host, self.port))
            s.listen(5)
            self.listen_socket = s
            
            while self.running:
                try:
                    conn, addr = s.accept()
                    if is_ip_compliant(addr[0], self.config['allowable_ip']):
                        serve = Server(conn, addr[0], addr[1])
                        serve.set_processor({
                            self.core: [MessageType.Register, MessageType.GetTask, MessageType.TaskResult, MessageType.Exit],
                            self._check_heartbeats: [MessageType.Heartbeat],
                        })
                        threading.Thread(
                            target=serve.start,
                        ).start()
                    else:
                        MessageProtocol.send(conn, {
                            'type': MessageType.Refuse,
                            'message': 'This address is not allowed to link'})
                        conn.close()
                        logging.info(f"This address [{addr}] is not allowed to link")
                except Exception as e:
                    logging.info(f"Server stop ... ...: {e}")
                    self.stop()
                    break

    def handle_register(self, msg: dict):
        if msg.get('type') == MessageType.Register:
            client_id, addr = msg['client_id'], msg['addr']
            logger.info(f"start register [{client_id}] [{addr}]")

            is_rgt = False  # 是否注册成功
            with self.lock:
                if client_id not in self.clients:
                    self.clients[addr] = time.time()
                    is_rgt = True
            if is_rgt:  # 注册成功
                logger.info(f"Slave host [{client_id}: {addr}] registered")
                self.clients_info[client_id] = {
                    'on_task': False,
                    'ip': addr[0],
                    'port': addr[1],
                }
                return True
        logger.info(f"This address [{addr}] is already registered")
        return False

    def core(self, serve: Server):
        if serve.running.is_set():
            # 注册
            msg = serve.get_message()
            serve.remote_name = msg['client_id']
            if self.handle_register(msg): 
                serve.send_message(build_allow())
            else: 
                serve.send_message(build_refuse('This address is already registered'))
                serve.stop()
                return
            
        while self.running.is_set():
            msg = serve.get_message()
            if msg.get('type') == MessageType.GetTask:
                serve.send_message(build_task({'id': '1', 'task': 'test'}))
            elif msg.get('type') == MessageType.TaskResult:
                serve.send_message(build_allow())
            


    def process_task(self, serve: Server):
        # 消息路由
        conn = serve.socket
        addr = serve.addr
        client_id = None
        while self.running:
            msg = serve.get_message()
            msg_type = msg.get('type')
            if msg_type == MessageType.Register:
                logging.info(f"start register [{addr}]")
                client_id = msg['client_id']
                is_rgt = False  # 是否注册成功
                with self.lock:
                    if client_id not in self.clients:
                        self.clients[client_id] = (conn, addr, time.time())
                        is_rgt = True
                if is_rgt:  # 注册成功
                    logging.info(f"Slave host [{client_id}: {addr}] registered")
                    self.clients_info[client_id] = {
                        'on_task': False,
                        'ip': addr[0],
                        'port': addr[1],
                    }
                    if not self.task_queue.empty():  # 任务不为空需要通知领取任务
                        serve.send_message({'type': MessageType.NewTasks, 'message': 'None'})
                else:  # 注册失败
                    serve.send_message({'type': MessageType.Refuse, 'message': 'This address is already registered'})
                
            elif msg_type == MessageType.Heartbeat:
                with self.lock:
                    if client_id in self.clients:
                        self.clients[client_id] = (
                            conn, addr, time.time()
                        )
                        
            elif msg_type == MessageType.GetTask:
                task = self._get_next_task()
                response = {'type': MessageType.Task, 'task': task}
                serve.send_message(response)
                
            elif msg_type == MessageType.TaskResult:
                self._handle_task_result(msg)
                
            elif msg_type == MessageType.SlaveInfo:
                self.update_slave_info()
                
            elif msg_type == MessageType.Exit:
                # 检查是否还有未完成任务
                # 安全退出
                serve.send_message({'type': MessageType.Allow})
                logging.info(f"[{addr}] Host exit safely")
                break
                # 拒绝退出
                
            elif msg_type == MessageType.ControlCommand:
                command = msg.get('command')
                if command == 'add_tasks':
                    tasks = msg.get('tasks', [])
                    logging.info(f"[{client_id}: {addr}] Control client added tasks: {tasks}")
                    self.dispatch_tasks(tasks)
                elif command == 'stop_server':
                    logging.info(f"[{client_id}: {addr}] Control client requested to stop the server")
                    self.stop()
                elif command == 'get_slave_info':
                    with self.lock:
                        slave_info = []
                        for id, (addr, _, _) in self.clients:
                            slave_info.append({id: addr})
                    serve.send_message({'type': MessageType.SlaveInfo, 'data': slave_info})
                    logging.info(f"[client: {addr}] Sends slave host information to the client")
                # 可以根据需要扩展更多命令

    def _get_next_task(self):
        """获取任务（生产者-消费者核心）"""
        if not self.task_queue.empty():
            return self.task_queue.get()
        return None

    def _handle_task_result(self, result: dict):
        """处理任务结果"""
        if not result['status']:
            self.task_queue.put(result['task'])  # 失败重试
            logging.warning(f"Task {result['task_id']} failed")
        else:
            logging.info(f"Task {result['task_id']} success")

    def _check_heartbeats(self, serve: Server):
        """心跳检测"""
        logging.info("Enable heartbeat detection")
        while self.running:
            try:
                time.sleep(HEARTBEAT_CHECK_PER)
                now = time.time()
                with self.lock:
                    to_remove = []
                    for cid, (_, _, last) in self.clients.items():
                        if now - last > HEARTBEAT_TIMEOUT:
                            to_remove.append(cid)
                    for cid in to_remove:
                        # TODO 检查是否还有未完成的任务
                        del self.clients[cid]
                        logging.warning(f"Client {cid} heartbeat detection timeout")
            except Exception as e:
                logging.error(f"Heartbeat detection error: {e}")
        logging.info(f"Heartbeat detection stopped")

    def _add_task(self, task: dict):
        """添加任务（线程安全）"""
        self.task_queue.put(task)

    def stop(self):
        if self.running:
            """停止服务"""
            self.running = False
            time.sleep(5)
            self.listen_socket.close()
        # logging.info("Master server stopped")

    def dispatch_tasks(self, tasks: list):
        """添加任务列表并给所有从服务器发送消息通知"""
        for task in tasks:
            self._add_task(task)
        logging.info(f"Task list added: {tasks}")
        send_info = []
        with self.lock:
            for id, (conn, addr, _) in self.clients.items():
                msg = {'type': MessageType.NewTasks, 'message': 'None'}
                MessageProtocol.send(conn, msg)
                send_info.append({
                    'id': id,
                    'addr': addr
                })
        logging.info(f"Sends a message of a new task to the servers: {send_info}")
            
    def update_slave_info():
        pass

def signal_handler(sig, frame):
    print("Interrupt signal received. Stopping servers...")
    if master: master.stop()
    sys.exit(0)
    
def server_start():
    # 启动主服务端
    global master
    master = MasterServer()
    master.start()

# ==================== 使用示例 ====================
if __name__ == '__main__':
    # 注册信号处理函数
    from common import TaskStatus
    t = TaskStatus(3)
    print(t)
    print(t.value)
    print(t.name)
